package ai.george.quickstart.controller;

import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;
import io.swagger.v3.oas.annotations.tags.Tag;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;


@Tag(name = "SSE示例")
@RestController
@RequestMapping("/sse")
public class SseController {


    private final ChatModel chatModel;

    private final ChatClient chatClient;


    public SseController(ChatModel chatModel) {
        this.chatModel = chatModel;
        this.chatClient = ChatClient.builder(chatModel).build();
    }


    @GetMapping(value = "/text1", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter text1(@RequestParam(name = "query", defaultValue = "你好，很高兴认识你，能简单介绍一下自己吗？") String query) {
        SseEmitter emitter = new SseEmitter();
        Message message = new UserMessage(query);
        Prompt prompt = new Prompt(message, ChatOptions.builder().model(DashScopeApi.ChatModel.QWEN_MAX.getValue()).maxTokens(2048).build());
        Flux<ChatResponse> flux = chatModel.stream(prompt);


        // 处理流式数据
        flux.doOnNext(chunk -> {
                    //这里的chunk就是每次返回的数据块
                    try {
                        String text = chunk.getResult().getOutput().getText();
                        //发送响应的文本给前端
                        emitter.send(text);
                    } catch (IOException e) {
                        //发生异常，发送错误到前端
                        emitter.completeWithError(e);
                    }
                })
                //发生异常，发送错误到前端
                .doOnError(e->{
                    emitter.completeWithError(e);
                })
                //流式数据处理完成，发送完成事件给前端
                .doOnComplete(()->{
                    emitter.complete();
                })
                //订阅流式数据
                .subscribe();
//                //发生异常，发送错误到前端
//                .doOnError(emitter::completeWithError)
//                //流式数据处理完成，发送完成事件给前端
//                .doOnComplete(emitter::complete)
//                //订阅流式数据
//                .subscribe();
        return emitter;
    }


    @GetMapping(value = "/text2", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter text2(@RequestParam(name = "query", defaultValue = "你好，很高兴认识你，你能帮我做些什么？") String query) {
        SseEmitter emitter = new SseEmitter();
        Message message = new UserMessage(query);
        Prompt prompt = new Prompt(message, ChatOptions.builder().model(DashScopeApi.ChatModel.QWEN_MAX.getValue()).maxTokens(2048).build());

        Flux<ChatResponse> flux = chatClient.prompt(prompt)
                .stream()
                .chatResponse();

        // 处理流式数据
        flux.doOnNext(chunk -> {
                    //这里的chunk就是每次返回的数据块
                    try {
                        String text = chunk.getResult().getOutput().getText();
                        //发送响应的文本给前端
                        emitter.send(text);
                    } catch (IOException e) {
                        //发生异常，发送错误到前端
                        emitter.completeWithError(e);
                    }
                })
                //发生异常，发送错误到前端
                .doOnError(emitter::completeWithError)
                //流式数据处理完成，发送完成事件给前端
                .doOnComplete(emitter::complete)
                //订阅流式数据
                .subscribe();
        return emitter;
    }
}
